# coding: utf8
'''
Created on Jun 21, 2011

@author: Nam Khanh Tran, Ba Dat Nguyen
'''
import os
import re
import math

import matplotlib.pyplot as plt

dict_path = "20-ng/train" 

def read_file(dict_path):
    """
    """
    vocab = dict()
    categories = dict()
    vocab_cat = dict()
    documents = dict()
    
    print 'Reading file....'
    for cat in os.listdir(dict_path):
        categories[cat] = list()
        print dict_path + "/" + cat
        for doc in os.listdir(dict_path + "/" + cat):            
            categories[cat].append(doc)
            with open(dict_path + "/" + cat + "/" + doc, 'r') as fin:
                content = fin.read()
                content = re.sub('[“„‘~!@#$%^&*(){}:;",\'.?-_=|\n><]', ' ', content.lower())
                word_list = re.split('\s+', content.strip())
                
                temp = list()
                for word in word_list:
                    if len(word) < 1 or word in temp:
                        continue
                    temp.append(word)
                    # word
                    if vocab.has_key(word):
                        vocab[word].append(doc)
                    else:
                        vocab[word] = [doc]
                    #word_cat
                    if vocab_cat.has_key((word,cat)):
                        vocab_cat[(word,cat)].append(doc)
                    else:
                        vocab_cat[(word,cat)] = [doc]
                #end for
                documents[doc] = temp;
            #end with
        #end for
    #end for
    return (vocab, categories, vocab_cat, documents)

def info_gain(vocab, categories, vocab_cat):
    """
    """
    print "\nComputing information gain...."
    
    number_documents = 0
    for k, doc_l in categories.items():
        number_documents = number_documents + len(doc_l)
        
    total_entropy = 0.0
    for k, doc_l in categories.items():
        prob_c = len(doc_l) * 1.0 / number_documents
        total_entropy = total_entropy - (prob_c * math.log(prob_c, 2))
    
    info_gain = dict()
    
    for word in vocab.keys():        
        prob_w = len(vocab[word]) * 1.0 / number_documents
        first_entropy = 0.0
        second_entropy = 0.0
        for cat in categories.keys():
            if vocab_cat.has_key((word,cat)):
                prob_w_c = len(vocab_cat[(word,cat)]) * 1.0 / len(vocab[word])
                first_entropy = first_entropy - prob_w_c * math.log(prob_w_c, 2)
                
                if number_documents == len(vocab[word]):
                    continue
                prob_nw_c = (len(categories[cat]) - len(vocab_cat[(word,cat)])) * 1.0 / (number_documents - len(vocab[word]))
                second_entropy = second_entropy - prob_nw_c * math.log(prob_nw_c, 2)
            else:
                if number_documents - len(vocab[word]) == 0:
                    continue
                prob_nw_c = len(categories[cat]) * 1.0 / (number_documents - len(vocab[word]))
                second_entropy = second_entropy - prob_nw_c * math.log(prob_nw_c, 2)
            
            info_gain[word] = total_entropy - prob_w * first_entropy - (1 - prob_w) * second_entropy

    pred_list = sorted(info_gain.items(), key=lambda(k,v):v, reverse=True)
    with open('ig_top100.txt', 'w') as fout:
        for k,v in pred_list[:100]:
            fout.write(k)
            fout.write('\n')
            
    with open('ig_top1000.txt', 'w') as fout:
        for k,v in pred_list[:1000]:
            fout.write(k)
            fout.write('\n')
                            
    info_gain = dict()
    for k,v in pred_list:
        info_gain[k] = v
        
    return info_gain

def chi_square(vocab, categories, vocab_cat):
    """
    """
    print "\nComputing chi-square..."
    
    chi_square = dict()
    
    number_documents = 0
    for k, doc_l in categories.items():
        number_documents = number_documents + len(doc_l)

    for word in vocab.keys():
        chi_square[word] = 0.0
        for cat in categories.keys():
            temp = 0.0
            if vocab_cat.has_key((word,cat)):
                f1 = len(vocab[word]) * len(categories[cat]) * 1.0 / number_documents
                f2 = len(vocab[word]) * (number_documents - len(categories[cat])) * 1.0 / number_documents
                f3 = (number_documents - len(vocab[word])) * len(categories[cat]) * 1.0 / number_documents
                f4 = (number_documents - len(vocab[word])) * (number_documents - len(categories[cat])) * 1.0 / number_documents
                
                temp += math.pow(len(vocab_cat[(word,cat)]) - f1, 2.0) / f1
                if f2 != 0:
                    temp += math.pow(len(vocab[word]) - len(vocab_cat[(word,cat)]) - f2, 2.0) / f2
                if f3 != 0:
                    temp += math.pow(len(categories[cat]) - len(vocab_cat[(word,cat)]) - f3, 2.0) / f3
                if f4 != 0:
                    temp += math.pow(number_documents - len(vocab[word]) - len(categories[cat]) + len(vocab_cat[(word,cat)]) - f4, 2.0) / f4

            chi_square[word] += (len(categories[cat]) * 1.0 / number_documents) * temp
            
    pred_list = sorted(chi_square.items(), key=lambda(k,v):v, reverse=True)
    with open('cs_top100.txt', 'w') as fout:
        for k,v in pred_list[:100]:
            fout.write(k)
            fout.write('\n')
            
    with open('cs_top1000.txt', 'w') as fout:
        for k,v in pred_list[:1000]:
            fout.write(k)
            fout.write('\n')
    
    chi_square = dict()
    for k,v in pred_list:
        chi_square[k] = v
        
    return chi_square
   
    
if __name__ == "__main__":
    (vocab, categories, vocab_cat, documents) = read_file(dict_path)
    with open('training_vocab.txt', 'w') as fout:
        for word in vocab.keys():
            fout.write(word)
            fout.write('\n')
            
    info_gain = info_gain(vocab, categories, vocab_cat)
    chi_square = chi_square(vocab, categories, vocab_cat)
    
    x = list()
    y = list()
    for k in info_gain.keys():
        y.append(info_gain[k])
        x.append(chi_square[k])
        
    plt.plot(x, y, '+')
    plt.show()                        